from datetime import datetime
from typing import List, Optional

from pydantic import BaseModel, EmailStr, Field
from enum import Enum, IntEnum


class StrEnum(str,Enum):
    pass


class IntelligenceModelTaskType(StrEnum):
    OCR = "ocr"
    CLASSIFY = "classify"
    DETECT = "detect"
    OBB = "obb"
    SEG = "seg"
    POSE = "pose"
    VIDEO_DETECT = "video_detect"
    VIDEO_CLASSIFY = "video_classify"

class DataTypeEnum(StrEnum):
    IMAGE = "image"
    VIDEO = "video"
    AUDIO = "audio"

class DataPartitionEnum(IntEnum):
    TRAIN = 1
    VAL = 2
    TEST = 3

class TaskStatusEnum(IntEnum):
    OPEN = 0  # 未查阅
    CLOSE = 1  # 标注
    DELIVER = 2  # 交付，完成划分数据集

class AnotationStatusEnum(StrEnum):
    UNASSIGNED = "未分配"
    ASSIGNED = "已分配"
    VIEWED = "已经浏览"  # 
    COMPLETED = "已标注"
    REVIEWED = "检查完成"
    ERROR = "检查错误"
    


class BaseAIModel(BaseModel):
    id: int
    name: Optional[str] = None
    frame: Optional[str] = None
    task: Optional[str] = None
    labels: Optional[dict] = None
    is_quad: Optional[bool] = False
    created_at: Optional[datetime]
    updated_at: Optional[datetime]


class ModelCreate(BaseModel):
    name: str = Field(example="行驶证外框四点模型")
    frame: str = Field(example="YOLOv8")
    task: IntelligenceModelTaskType = Field(example=IntelligenceModelTaskType.POSE)
    labels: dict = Field(example={0:{"name":'行驶证主证正面',"color":'red',},1:{"name":'行驶证副证正面',"color":'red',},2:{"name":'行驶证主证背面',"color":'red',},})
    is_quad: Optional[bool] = True

    def create_dict(self):
        return self.model_dump(exclude_unset=True, exclude={})


class ModelUpdate(BaseModel):
    id: int
    name: str
    frame: str
    task: str
    labels: Optional[dict]
    is_quad: Optional[bool] = False

